# ready the dataset, Not use ImageFolder as the author did
import linecache
import random
import PIL.ImageOps
import PIL
from PIL import Image
import numpy as np
import torch
from torch.utils.data import Dataset

from Siamese.Config import Config


class OneDataset(Dataset):
    def __init__(self, txt, transform=None, target_transform=None, should_invert=False,
                 onlysame=False):
        self.transform = transform
        self.target_transform = target_transform
        self.should_invert = should_invert
        self.onlysame = onlysame
        self.pairs = []
        ori = open(txt, 'r')
        lines = ori.readlines()

        for line in lines:
            label = line.split()[1]
            if int(label) not in Config.typelist:
                continue
            self.pairs.append(line + ' ' + label)

        print(len(self.pairs))

    def __getitem__(self, index):
        pair = self.pairs[index].split()
        img0 = Image.open(pair[0])
        # img0 = img0.convert("L")

        if self.should_invert:
            img0 = PIL.ImageOps.invert(img0)

        if self.transform is not None:
            img0 = self.transform(img0)

        return img0, torch.from_numpy(np.array([int(pair[1])], dtype=np.float32))

    def __len__(self):
        return len(self.pairs)
